#!/usr/bin/env python
# encoding=utf-8
#codeby     道长且阻
#email      ydhcui@suliu.net/QQ664284092
#website    http://www.suliu.net


import time
import uuid
import hashlib
import json

from lib import sql
from settings import DATABASE

if DATABASE['datatype']=='mysql':
    userdata = MySQLDatabase(
        host     = DATABASE['datahost'],
        user     = DATABASE['username'],
        passwd   = DATABASE['password'],
        database = DATABASE['dataname'],
        port     = DATABASE['dataport'],
        charset  = DATABASE['charset'])
else:
    userdata = sql.SqliteDatabase(
        DATABASE['dataname']) # threadlocals = True
userdata.connect()

class BaseModel(sql.Model):
    class Meta:
        database = userdata

###########################################################
class Project(BaseModel):
    pj_id       = sql.CharField(unique = True, default = uuid.uuid4)
    pj_name     = sql.CharField()
    pj_desc     = sql.CharField(default = '')
    reportpath  = sql.CharField(default = '')
    createdate  = sql.IntegerField(default = time.time)
    finishdate  = sql.IntegerField(default = 0)

class DictResult(BaseModel):
    project      = sql.ForeignKeyField(Project)
    key          = sql.CharField()
    value        = sql.CharField()
    style        = sql.CharField(default = 'ALL')
    count        = sql.IntegerField(default = 1) #出现次数

class ScanTask(BaseModel):
    """扫描任务表"""
    project     = sql.ForeignKeyField(Project)
    task_id     = sql.CharField(unique = True, default = uuid.uuid4)
    task_name   = sql.CharField()                               #任务名称-对应task中的名称
    task_host   = sql.CharField()                               #任务主机
    task_args   = sql.CharField(default = '')                   #任务参数
    task_desc   = sql.CharField(default = '')                   #任务描述
    task_note   = sql.CharField(default = '')                   #备注
    task_pid    = sql.CharField(default = '')                   #任务对应的进程id
    task_state  = sql.CharField(default = 'runing')             #任务状态
    task_level  = sql.CharField(default = '0')                  #任务等级 0、1、2
    createdate  = sql.IntegerField(default = time.time)  #创建时间
    finishdate  = sql.IntegerField(default = 0)                 #0表示未完成

    def _update_task(self,task_pid,task_host,task_args=''):
        task,cd = self.get_or_create(
            task_pid = task_pid,
            task_host = task_host,
            task_args = task_args)
        task.createdate = time.time
        task.task_state = 'runing'
        task.save()
        return task,cd

    def _update_status(self,state='runing'):
        self.state      = state
        self.finishdate = time.time
        self.save()


class VulResult(BaseModel):
    '''漏洞扫描结果'''
    task         = sql.ForeignKeyField(ScanTask)
    name         = sql.CharField()
    address      = sql.CharField() #如果是http则是url，否则就是IP地址:端口
    request      = sql.CharField(default = '')
    response     = sql.CharField(default = '')
    createdate   = sql.IntegerField(default = time.time)

class HostResult(BaseModel):
    host         = sql.CharField()
    hostname     = sql.CharField(default = '')      #主机名
    version      = sql.CharField(default = '')      #系统版本

class PortResult(BaseModel):
    """端口扫描结果"""
    task         = sql.ForeignKeyField(ScanTask)
    host         = sql.ForeignKeyField(HostResult)
    port         = sql.IntegerField()
    type         = sql.CharField(default = '')       #tcp/udp
    name         = sql.CharField(default = '')       #http
    state        = sql.CharField(default = '')       #open
    product      = sql.CharField(default = '')       #
    extrainfo    = sql.CharField(default = '')
    reason       = sql.CharField(default = '')
    version      = sql.CharField(default = '')   #一般是指中间件版本
    cpe          = sql.CharField(default = '')
    createdate   = sql.IntegerField(default = time.time)

class HttpResult(BaseModel):
    """域名扫描结果表"""
    task         = sql.ForeignKeyField(ScanTask)
    host         = sql.ForeignKeyField(HostResult)
    port         = sql.IntegerField(default = 80)
    state        = sql.IntegerField(default = 200)
    domain       = sql.CharField(default = '')
    banner       = sql.CharField(default = '')
    xpoweredby   = sql.CharField(default = '')
    title        = sql.CharField(default = '')
    headers      = sql.TextField(default = '')
    content      = sql.TextField(default = '')
    createdate   = sql.IntegerField(default = time.time)


###########################################################################################
class User(BaseModel):
    uid          = sql.PrimaryKeyField()
    uname        = sql.CharField(max_length = 20, default='')                     #性名
    upwd         = sql.CharField(max_length = 32, null = False)                   #密码
    uemail       = sql.CharField(max_length = 20, default='')                     #邮箱
    createdate   = sql.IntegerField(default = time.time)                   #创建时间
    lastlogin    = sql.IntegerField(default = time.time)                   #最后登录时间
    token        = sql.CharField(max_length = 32, default='')                     #对应models中api的token 权限

    @staticmethod
    def _create_password(raw):
        return hashlib.sha256(raw.encode()).hexdigest().hexdigest() 

    def _check_password(self, raw):
        return self.create_password(raw) == self.upwd

    def _update_login(self):
        self.lastlogin = time.time()
        self.save()
##############################################################################################################
def main():
    [name.create_table(fail_silently=True) \
        for name in BaseModel.__subclasses__() \
            if not name.__name__.startswith('Base')]

if __name__ == '__main__':
    main()
    #print(Api().do({'model':'util','action':'gethostbyname','kwargs':{'name':'127.0.0.1'}}))

